#include <iostream>
#include <fstream>
#include <string>
#include <vector>
using namespace std;

// S = 32K, L = 258, BLOCK_SIZE = arbitrary
#include "deflate.h"
#include "hufftree.h"

int LZ77_Encode(const unsigned char *T, unsigned int length, vector<Triplet> &triplets, const int S, const int L, const int MIN_MATCH)
{
    unsigned short seq_len, best_len, best_offset;
    const unsigned char *pos = T, *end = T + length, *haystack, *needle;
    while(pos < end)
    {
        haystack = pos-1;
        needle = pos;
        best_len = best_offset = 0;
        while((haystack >= T) && ((pos - haystack) < S))
        {
            for(seq_len = 0; (haystack[seq_len] == needle[seq_len]) && (seq_len < L) && ((needle+seq_len+1) < end); ++seq_len);
            if((seq_len > MIN_MATCH) && (seq_len > best_len)) { best_len = seq_len; best_offset = pos - haystack; }
            --haystack;
        }
        pos += best_len;
        triplets.push_back(Triplet(best_offset, best_len, (*pos)));
        ++pos;
    }
    return triplets.size();
}

int LZ77_Decode(const vector<Triplet> &triplets, unsigned char *T)
{
    unsigned short len;
    unsigned char *pos = T, *ptr;
    for(size_t i = 0, im = triplets.size(); i < im; ++i, ++pos)
    {
        for(ptr = pos - triplets[i].i, len = triplets[i].j; len > 0; --len, ++ptr, ++pos)
            (*pos) = (*ptr);
        (*pos) = triplets[i].a;
    }
    (*pos) = 0;
    return pos - T; // length
}

// looks at `pos`th bit in the `stream` (starting from MSB)
unsigned char BitStream_BitView(const unsigned char *stream, unsigned int pos)
{
    unsigned char byte = stream[pos/8];
    byte >>= (7 - (pos % 8));
    return (unsigned char)(byte % 2);
}

// bits - size in bits
unsigned int BitStream_Bytes(unsigned int bits)
{
    return ((bits / 8) + ((bits % 8) != 0));
}

// write_bits[in] - number of bits counted from LSB (because of type casting) in range 1-16, because of val's data type (2 byte short)
// val[in] - value that I want to add into the `stream`, pos[in|out] - what bit am I at inside of `encoded`, encoded[in|out] - encoded bit stream
void BitStream_Add(int write_bits, unsigned short val, unsigned char *stream, unsigned int &pos)
{
    if((write_bits < 1) || (write_bits > 16)) return;
    //
    unsigned int byte_pos;
    unsigned int wbits;
    unsigned char space;
    unsigned char bytes[2];
    unsigned char nbits[2];
    if(write_bits > 8)
    {
        bytes[0] = (unsigned char)(val >> 8); bytes[1] = (unsigned char)val;
        nbits[0] = (write_bits - 8);          nbits[1] = 8;
    }
    else
    {
        bytes[0] = (unsigned char)val; bytes[1] = 0;
        nbits[0] = write_bits;         nbits[1] = 0;
    }
    //
    for(int b = 0; (b < 2) && (nbits[b] > 0); b++)
    {
        byte_pos = pos / 8;   // array index to stream
        //
        space = 8 - nbits[b];	// reuse the variable `space` for shift parameter
        bytes[b] = ((bytes[b] << space) >> space);	// shifts are there to make sure that bits higher than `nbits` are all 0 (it could be done by mask but it would had to be (pre)calculated)
        //
        wbits = pos % 8;    // number of bits written in the last byte
        space = 8 - wbits;	// space left in the last byte
        if(nbits[b] <= space)
        {
            stream[byte_pos] |= (bytes[b] << (space - nbits[b]));
            pos += nbits[b];
        }
        else //if(nbits[b] > space)
        {
            wbits = nbits[b] - space;   // number of bits that are over the limit
            stream[byte_pos++] |= (bytes[b] >> wbits);	// remove bits that will not fit inside of the byte and add the rest
            stream[byte_pos  ] |= (bytes[b] << (8 - wbits)); // add the previously removed part so it'll be complete
            pos += nbits[b];
        }
    }
}

unsigned short byte_swap(unsigned short word)
{
	unsigned char byte = (word >> 8);	// store higher bits
	word <<= 8;	// shift lower bits to higher positions
	word |= byte;	// finish the swap
	return word;
}

unsigned short bit_reverse(unsigned short bit_len, unsigned short word)	// data-type short is used only for easier bit operations (argument should be of type unsigned char[2])
{
	if(bit_len > 16) return 0;
	//
	// Warning:
	// ========
	// #1) ENDIAN! bytes in numerical data types are backwards
	// #2) BitVeiw stream must have all bits left side (MSB) aligned - bit 0 is MSB, bit 7 is LSB
	// ========
	// -> 1) shift bits to the MSB position (because of #2)
	word <<= 16 - bit_len;
	// -> 2) swap bytes (because of #1)
	word = byte_swap(word);
	// -> 3) get the reversal
    unsigned short res = 0;
    for(unsigned short b = 0; b < bit_len; ++b)
        res |= BitStream_BitView((const unsigned char *)&word, b) << b;
    return res;
}

// symbol[in] - symbol to be decoded, pos[in|out] - what bit am I at inside of `encoded`, encoded[in|out] - encoded bit stream
void Huffman_Encode_Symbol(unsigned short symbol, unsigned int &pos, unsigned char *encoded)
{	// bits must be inseted in reverse order to be properly decoded!!
    BitStream_Add(static_ltree_lengths[symbol], bit_reverse(static_ltree_lengths[symbol], static_ltree_encoder[symbol]), encoded, pos);
}

// encoded[in] - encoded bit stream, pos[in|out] - what bit am I at inside of `encoded`, symbol[out] = 1 decoded symbol
void Huffman_Decode_Symbol(const unsigned char *encoded, unsigned int &pos, unsigned short &symbol)
{	// this algorithm was taken over from gziplite-3.0
    int len;            /* current number of bits in code */
    int code;           /* len bits being decoded */
    int first;          /* first code of length len */
    int count;          /* number of codes of length len */
    int index;          /* index of first code of length len in symbol table */
	code = first = index = 0;
    for (len = 1; len <= 9; len++)	// 9 is maximal number of bits per codeword
	{
        code |= BitStream_BitView(encoded, pos++);   /* get next bit */
	    count = bl_count[len];
        if(code - count < first)       /* if length len, return symbol */
		{
            symbol = static_ltree_decoder[index + (code - first)];
			break;
		}
        index += count;                 /* else update for next length */
        first += count;
        first <<= 1;
        code <<= 1;
    }
}

void Distance_Encode(unsigned short val, unsigned char *encoded, unsigned int &pos)
{	// this might not be compatible with the deflate algorithm, because bits might be inseted in reverse order (I'm not sure)
	//assert(val == 0)	<-- this vill never happen, because there can't be distance of 0!
    val -= 1;
    unsigned short bi; // distance base index (0-29)
    for(bi = 0; bi < 30; bi++)
        if(base_dist[bi] > val)
            break;
    //
    bi -= 1;
    //
    BitStream_Add(5, static_dtree_encoder[bi], encoded, pos);
    if(bi > 3) BitStream_Add((bi/2)-1, val-base_dist[bi], encoded, pos);
}

void Distance_Decode(const unsigned char *encoded, unsigned int &pos, unsigned short &val)
{	// this might not be compatible with the deflate algorithm, because bits might be inseted in reverse order (I'm not sure)
    unsigned short bi = 0; // distance base index (0-29)
    for(int i = 0; i < 5; ++i, ++pos)
        bi = (bi << 1) | BitStream_BitView(encoded, pos);
    bi = static_dtree_decoder[bi];
    //
    unsigned short off = 0; // distance offset
    for(int i = 0, im = (bi/2)-1; i < im; ++i, ++pos)
        off = (off << 1) | BitStream_BitView(encoded, pos);
    //
    val = base_dist[bi] + off + 1;
}

void Length_Encode(unsigned short value, unsigned char *encoded, unsigned int &pos)
{	// this might not be compatible with the deflate algorithm, because extra bits might be inseted in reverse order (I'm not sure)
    if((value < 3) || (value > 258)) return;
    //
    value -= 3;
    int bi;
    for(bi = 0; bi < 29; ++bi)
        if(base_length[bi] > value)
            break;
    //
    bi -= 1;
    Huffman_Encode_Symbol(257+bi, pos, encoded);
    if((bi > 7) && (bi < 28)) BitStream_Add((bi/4)-1, value - base_length[bi], encoded, pos); // extra bits, if any
}

void Length_Decode(unsigned short lenc, unsigned char *encoded, unsigned int &pos, unsigned short &value)
{	// this might not be compatible with the deflate algorithm, because extra bits might be inseted in reverse order (I'm not sure)
    int bi = lenc - 256 - 1;
    unsigned short lextra = 0;
    for(int i = 0, im = (bi/4)-1; i < im; ++i, ++pos)
        lextra = (lextra << 1) | BitStream_BitView(encoded, pos);
    //
    value = 3 + base_length[bi] + lextra;
}

// inserts END OF BLOCK symbol (code 256)
void Insert_EndOfBlock_Symbol(unsigned int &pos, unsigned char *encoded)
{
	Huffman_Encode_Symbol(256, pos, encoded);
}